#' 估计样本重叠导致的偏倚
#'
#' 代码改编自https://github.com/mglev1n/mrSampleOverlap，为了避免依赖的R包太多，
#' 这里将其移植并改编到这里并适应本包的数据。
#' 原始文献为https://onlinelibrary.wiley.com/doi/10.1002/gepi.21998
#'
#' @param sub_res res的单暴露和单结局的亚集
#' @param overlap_prop 样本重叠率，假设暴露是UKB(400000人)+FinnGen(200000人)的600000人GWAS meta分析后的数据，
#' 结局是UKB(400000人)+PGC(300000人)的700000人GWAS meta分析后的数据，那么样本重叠率为400000/700000，
#' 分母取较大的数，也就是700000。
#'
#' @param ols_bias Bias in Observational Estimate，有人说问了作者就是IVW的beta值，但是还没有看到
#' 有文献这么做过，所以目前还是个疑问
#'
#' @return 图或者数字
#' @export
#'
#' @examples
#'
#' \dontrun{
#'
#' # 假设暴露'ieu-a-300'和结局'ieu-a-297'的样本存在交集，重叠率为0.1,ols_bias为0.3
#' sub_res<-subset( res,id.exposure %in% 'ieu-a-300' & id.outcome %in% 'ieu-a-297' )
#' bias <- U8_estimate_overlap_bias(sub_res,0.1,0.3)
#' bias
#' # 如果重叠率和ols_bias都不填的话,会作一个图
#' bias <- U8_estimate_overlap_bias(sub_res)
#' bias
#' # 如果计算重叠率为0.1，只有ols_bias不填的话,也会作一个图
#' bias <- U8_estimate_overlap_bias(sub_res,0.1)
#' bias
#'
#' }
#'
#'
#'
#'
#'
U8_estimate_overlap_bias<-function(sub_res,
                                   overlap_prop=NULL,
                                   ols_bias=NULL){

  if( nrow(sub_res) >1 ){
    stop("只能用于单个暴露到单个结局")
  }

  if(sub_res$type.outcome=="Binary"){
    sub_res$case_prop<-  sub_res$ncase.outcome/sub_res$samplesize.outcome
  }else{
    sub_res$case_prop <-  0
  }

  if( is.null(ols_bias ) ){

    suppressPackageStartupMessages(require(tidyr))
    suppressPackageStartupMessages(require(dplyr))
    suppressPackageStartupMessages(require(ggplot2))

    if(is.null(overlap_prop)){

      grid <- crossing(overlap_prop = seq(0, 1, 0.1),
                       ols_bias = seq(0, 1, 0.2))
      bias_res <- sub_res %>%
        crossing(grid) %>%
        mutate(res = estimate_overlap_bias(samplesize_exposure=sub_res$samplesize.exposure,
                                           samplesize_outcome=sub_res$samplesize.outcome,
                                           n_variants=sub_res$nSNP,
                                           rsq_exposure = sub_res$total_R2,
                                           case_prop = sub_res$case_prop,
                                           ols_bias = ols_bias,
                                           overlap_prop = overlap_prop)) %>%
        unnest(res)

      bias<-bias_res %>%
        pivot_longer(cols = c(bias, type1_error)) %>%
        ggplot(aes(overlap_prop, value, group = ols_bias, color = as.character(ols_bias))) +
        geom_point() +
        geom_line() +
        facet_grid(rows = vars(exposure),
                   cols = vars(name),
                   scales = "free_y") +
        labs(x = "Proportion of Overlapping Participants",
             y = "Value") +
        scale_color_discrete(name = "Bias in \nObservational \nEstimate") +
        theme_bw(base_size = 14)


    }else{

      grid <- crossing(ols_bias = seq(0, 1, 0.2))
      bias_res <- sub_res %>%
        crossing(grid) %>%
        mutate(res = estimate_overlap_bias(samplesize_exposure=sub_res$samplesize.exposure,
                                           samplesize_outcome=sub_res$samplesize.outcome,
                                           n_variants=sub_res$nSNP,
                                           rsq_exposure = sub_res$total_R2,
                                           case_prop = sub_res$case_prop,
                                           ols_bias = ols_bias,
                                           overlap_prop = overlap_prop)) %>%
        unnest(res)%>%
        mutate(overlap_prop=overlap_prop  )

      bias<- bias_res %>%
        pivot_longer(cols = c(bias, type1_error)) %>%
        ggplot(aes(ols_bias, value, group = name, color = name )) +
        geom_point() +
        geom_line() +
        facet_grid(rows = vars(exposure),
                   cols = vars(name),
                   scales = "free_y") +
        labs(x = "Bias in Observational Estimate",
             y = "Value") +
        scale_color_discrete(name = "Bias") +
        theme_bw(base_size = 14)

    }


  }else{

    bias <- estimate_overlap_bias(samplesize_exposure=sub_res$samplesize.exposure,
                                  samplesize_outcome=sub_res$samplesize.outcome,
                                  n_variants=sub_res$nSNP,
                                  rsq_exposure = sub_res$total_R2,
                                  case_prop = sub_res$case_prop,
                                  ols_bias = ols_bias,
                                  overlap_prop = overlap_prop )
  }
  return(bias)

}

estimate_overlap_bias <- function(samplesize_exposure, samplesize_outcome, n_variants, rsq_exposure, exp_f = NULL, lci_95 = FALSE, case_prop = 0, ols_bias, overlap_prop, var_x = 1, var_y = 1) {
  # adapted from Burgess et. al. (2016) PMID: 27625185

  if (!is.null(exp_f)) {
    expf <- exp_f
  } else {
    expf <- estimate_f(samplesize_exposure, n_variants, rsq_exposure, lci_95)
  }

  if (case_prop == 0) {
    var <- var_y / (samplesize_outcome * var_x * rsq_exposure)
  } else {
    var <- 1 / (samplesize_outcome * rsq_exposure * var_x * case_prop * (1 - case_prop))
  }

  bias <- ols_bias * overlap_prop * (1 / expf)

  type1_error <- 2 - pnorm(1.96 + bias / sqrt(var)) - pnorm(1.96 - bias / sqrt(var))

  return(data.frame(bias = bias, type1_error = type1_error))
}

estimate_f <- function(samplesize_exposure, n_variants, rsq_exposure, lci_95 = FALSE) {
  expf <- (samplesize_exposure - n_variants - 1) / n_variants * rsq_exposure / (1 - rsq_exposure)

  if (lci_95) {
    expf <- estimate_f_lci95(expf, n_variants, samplesize_exposure)
  }

  return(expf)
}

estimate_f_lci95 <- function(f, nu1, nu2) {
  lambda <- f * nu1 * (nu2 - 2) / nu2 - nu1
  lower <- f - 1
  while (pf(lower, df1 = nu1, df2 = nu2, ncp = lambda) > 0.05) {
    lower <- lower - 1
  }
  upper <- lower + 1
  while (abs(pf((lower + upper) / 2, df1 = nu1, df2 = nu2, ncp = lambda) - 0.05) > 0.0001) {
    if (pf((lower + upper) / 2, df1 = nu1, df2 = nu2, ncp = lambda) > 0.05) {
      upper <- (lower + upper) / 2
    }
    if (pf((lower + upper) / 2, df1 = nu1, df2 = nu2, ncp = lambda) < 0.05) {
      lower <- (lower + upper) / 2
    }
  }
  return((lower + upper) / 2)
}
